import numpy as np
import cPickle
import os
import h5py

# input param:
# candis: bbox for each scan
# thre_dist : 10 mm
# spac: spacing of vol
# nms_para: if nms: 1 else: 0
# return param:
# det: center of nmsed candidates
# bbx: bbx of nmsed candidates
# prob: score of nmsed candidates



def nms(candis,score,thre_dist,spac,nms_para):
    if nms_para:
        center = []
        num_candis = len(candis)
        for icandis in xrange(num_candis):
            x=(candis[icandis][1]+candis[icandis][3])/2
            y=(candis[icandis][0]+candis[icandis][2])/2
            z=candis[icandis][4]
            center.append(np.hstack([x,y,z]))

        dists = np.zeros([len(candis),len(candis)])
        for r in xrange (len(candis)):
            for c in xrange(len(candis)):
                dists[r][c]=((center[r][0]-center[c][0])*spac[0])**2+ \
                ((center[r][1]-center[c][1])*spac[1])**2+((center[r][2]-center[c][2])*spac[2])**2

        order = sorted(range(len(score)),key=lambda k: score[k],reverse=True)
        keep=[]
        while len(order)>0:
            i =order[0]
            keep.append(i)
            inds=[]
            for j in xrange(len(order)):
                if  dists[i][order[j]]>thre_dist**2:
                    inds.append(j)



            tmp=[]
            for i_ind in xrange(len(inds)):
                tmp.append(order[inds[i_ind]])
            order=tmp


    else:
        keep=range(len(candis))

    det=[]
    bbx=[]
    prob=[]
    for idet in xrange(len(keep)):
        x=(candis[keep[idet]][1]+candis[keep[idet]][3])/2
        y=(candis[keep[idet]][0]+candis[keep[idet]][2])/2
        z=candis[keep[idet]][4]
        det.append(np.hstack([x,y,z]))
        bbx.append(np.hstack([candis[keep[idet]][1],candis[keep[idet]][0],candis[keep[idet]][4],candis[keep[idet]][3],candis[keep[idet]][2],candis[keep[idet]][4]+1]))
        prob.append(score[keep[idet]])
    return det,bbx,prob

def evaluate(center,radius,lab,spac):
    candisPos=[]
    candisNeg=[]
    isFound=np.zeros(len(lab))
    if len(lab)>0:
        for icandis in xrange(len(center)):
            flag=0
            for ilab in xrange(len(lab)):
                can=center[icandis]
                cen=lab[ilab]
                if ((can[0]-cen[0])*spac[0])**2+((can[1]-cen[1])*spac[1])**2+((can[2]-cen[2])*spac[2])**2<(2+radius[ilab])**2:
                    isFound[ilab]=1
                    candisPos.append(can)
                    flag=1
            if flag==0:
                candisNeg.append(center[icandis])
    else:
        candisNeg=center
    return len(candisNeg),len(candisPos),isFound


def main_nms_five_fold(candis_dir,lungmask_lidc_dir,lungmask_kaggle_dir,vol_lidc_dir,vol_kaggle_dir,thre_score,thre_dist,thre_xy):
    D=os.listdir(candis_dir)
    l1=np.load('/data-disk/Kaggle/backup/lidc/label_dict_3_0329.npy').item()
    l2=np.load('/data-disk/Kaggle/backup/kaggle/label_dict_0325.npy').item()

    f1=h5py.File(lungmask_lidc_dir,'r')
    f2=h5py.File(lungmask_kaggle_dir,'r')
    v1=h5py.File(vol_lidc_dir,'r')
    v2=h5py.File(vol_kaggle_dir,'r')

    num_pos_t=0
    num_neg_t=0
    find_nod=0
    find_ma=0
    all_nod=0
    all_ma=0
    candidates_list=[]
    for iscan in xrange(len(D)):
        foldName=D[iscan]
        candis_ori=np.load(candis_dir+foldName+'/candis.npy')
        if foldName[0]=='L':
            spac = v1[foldName].attrs['spacing']
            lungmask=f1[foldName]
            source_str='lidc'
            lab = l1[foldName]
            center=[]
            radius=[]
            malign=[]
            for ilab in xrange (len(lab)):
                center.append(lab[ilab]['center'])
                radius.append(lab[ilab]['radius'])
                malign.append(lab[ilab]['attrs']['malignancy'])

        else:
            spac = v2[foldName].attrs['spacing']
            lungmask=f2[foldName]
            source_str='kaggle_mali'
            lab = l2[foldName]
            center=[]
            radius=[]
            malign=[]
            for ilab in xrange (len(lab)):
                center.append(lab[ilab]['center'])
                radius.append(lab[ilab]['radius'])
                malign.append(lab[ilab]['malign'])



        candis=[]
        score=[]
        for islice in xrange(len(candis_ori)):
            tmp=candis_ori[islice]
            if len(tmp):
                for j in xrange(len(tmp)):
                    if tmp[j][4]<thre_score:
                        continue
                    if max(tmp[j][2]-tmp[j][0],tmp[j][3]-tmp[j][1])*spac[0]<thre_xy:
                        continue
                    if lungmask[(tmp[j][1]+tmp[j][3])/2,(tmp[j][0]+tmp[j][2])/2,islice]==0 and max(tmp[j][2]-tmp[j][0],tmp[j][3]-tmp[j][1])*spac[0]<=15:
                        continue
                    candis.append(np.hstack([tmp[j][0:4],islice]))
                    score.append(tmp[j][4])

        # print candis_ori
        candis_a,candis_bbx,prob_a=nms(candis,score,thre_dist,spac,1)
        for idict in xrange(len(candis_bbx)):
            candidate_dict=dict(scanid=foldName, source=source_str, bbox=candis_bbx[idict], center=candis_a[idict],prob=prob_a[idict])
            candidates_list.append(candidate_dict)

        # compute accu:  if u don't want to evaluate, ignore

        # num_neg,num_pos,isFound=evaluate(candis_a,radius,center,spac)

        # num_pos_t=num_pos_t+num_pos
        # num_neg_t=num_neg_t+num_neg
        # find_nod=find_nod+sum(isFound)
        # all_nod=all_nod+isFound.shape[0]

        # malign=np.array(malign)
        # if foldName[0]=='L':
        # 	find_ma=find_ma+sum(isFound*((malign>3)*1))
        # 	all_ma=all_ma+sum((malign>3)*1)
        # else:
        # 	find_ma=find_ma+sum(isFound*((malign>6)*1))
        # 	all_ma=all_ma+sum((malign>6)*1)

        # print iscan,foldName,find_nod,all_nod,find_ma,all_ma,num_pos_t/(1.0+iscan),num_neg_t/(1.0+iscan)

    return candidates_list


def main_nms(candis_dir,lungmask_dir,vol_dir,thre_score,thre_dist,thre_xy,source_str):
    # candis_dir='/data-disk/Kaggle/candidates/workingdir/0330/candis_kaggle_testset_before_nms_lungmask_60000/'
    D=os.listdir(candis_dir)
    # f2=h5py.File('/data-disk/Kaggle/backup/kaggle/lungmask.hdf5','r')
    # v2=h5py.File('/data-disk/Kaggle/backup/kaggle/vol.hdf5','r')
    f2=h5py.File(lungmask_dir,'r')
    v2=h5py.File(vol_dir,'r')


    num_nod=0
    candidates_list=[]
    for iscan in xrange(len(D)):
        foldName=D[iscan]
        candis_ori=np.load(candis_dir+foldName+'/candis.npy')
        lungmask=f2[foldName]
        # source_str='kaggle_test'
        spac = v2[foldName].attrs['spacing']


        candis=[]
        score=[]
        for islice in xrange(len(candis_ori)):
            tmp=candis_ori[islice]
            if len(tmp):
                for j in xrange(len(tmp)):
                    if tmp[j][4]<thre_score:
                        continue
                    if max(tmp[j][2]-tmp[j][0],tmp[j][3]-tmp[j][1])*spac[0]<thre_xy:
                        continue
                    if lungmask[(tmp[j][1]+tmp[j][3])/2,(tmp[j][0]+tmp[j][2])/2,islice]==0 and max(tmp[j][2]-tmp[j][0],tmp[j][3]-tmp[j][1])*spac[0]<=15:
                        continue
                    candis.append(np.hstack([tmp[j][0:4],islice]))
                    score.append(tmp[j][4])

        candis_a,candis_bbx,prob_a=nms(candis,score,thre_dist,spac,1)
        for idict in xrange(len(candis_bbx)):
            candidate_dict=dict(scanid=foldName, source=source_str, bbox=candis_bbx[idict], center=candis_a[idict],prob=prob_a[idict])
            candidates_list.append(candidate_dict)
        # num_nod=num_nod+len(candis_bbx)
        # print iscan,foldName,num_nod/(iscan+1.0)
    return candidates_list


def nms_one_scan(candis, thresh_v, )

if __name__ == '__main__':

    ## for kaggle test/ kaggle beni/spie
    candis_dir='/data-disk/Kaggle/candidates/workingdir/0330/candis_kaggle_testset_before_nms_lungmask_60000/'
    lungmask_dir='/data-disk/Kaggle/backup/kaggle/lungmask.hdf5'
    vol_dir='/data-disk/Kaggle/backup/kaggle/vol.hdf5'
    thre_xy= 5 # bbox>5mmm
    thre_dist =10 # 10mm
    thre_score =0.1 # this u can choose v9: 0.98 v15:0.1
    source_str='kaggle_testset' #source
    candidate_list=main_nms(candis_dir,lungmask_dir,vol_dir,thre_score,thre_dist,thre_xy,source_str)
    ## for five fold
    candis_dir='/data-disk/Kaggle/candidates/workingdir/0330/candis_5fold_before_nms_lungmask_60000/'
    lungmask_kaggle_dir='/data-disk/Kaggle/backup/kaggle/lungmask.hdf5'
    lungmask_lidc_dir='/data-disk/Kaggle/backup/lidc/lungmask.hdf5'
    vol_kaggle_dir='/data-disk/Kaggle/backup/kaggle/vol.hdf5'
    vol_lidc_dir='/data-disk/Kaggle/backup/lidc/vol.hdf5'
    thre_xy= 5 #bbox>5mm
    thre_dist =10 # 10mm
    thre_score =0.1 # this u can choose v9: 0.98 v15:0.1
    # source str is lidc and kaggle_magli, which has already been set in the program
    candidate_list=main_nms_five_fold(candis_dir,lungmask_lidc_dir,lungmask_kaggle_dir,vol_lidc_dir,vol_kaggle_dir,thre_score,thre_dist,thre_xy)

